package org.alis.rpc.manager;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioSocketChannel;
import lombok.extern.slf4j.Slf4j;
import org.alis.rpc.common.protocol.RpcProtocol;
import org.alis.rpc.common.protocol.RpcServiceInfo;
import org.alis.rpc.common.util.ThreadPoolUtils;
import org.alis.rpc.handler.RpcClientHandler;
import org.alis.rpc.handler.RpcClientInitializer;
import org.alis.rpc.route.AbstractRpcLoadBalance;
import org.alis.rpc.route.impl.RpcLoadBalanceRoundRobin;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * ConnectManager desc: 链接管理器 管理链接线程池
 *
 * @author luochuan
 * @date 2020/9/30
 */
@Slf4j
public class ConnectManager<V> {
    private EventLoopGroup eventLoopGroup = new NioEventLoopGroup(4);
    private Map<RpcProtocol, RpcClientHandler<V>> connectedServerNodes = Maps.newConcurrentMap();
    private Set<RpcProtocol> rpcProtocolSet = Sets.newCopyOnWriteArraySet();
    private ReentrantLock lock = new ReentrantLock();
    private Condition connected = lock.newCondition();
    private AbstractRpcLoadBalance loadBalance = new RpcLoadBalanceRoundRobin();
    private volatile boolean isRunning = true;

    private static ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(4, 8,
            600L, TimeUnit.SECONDS, new LinkedBlockingQueue<>(1000),ThreadPoolUtils.NameThreadFactory.build("client-manager"),
            new ThreadPoolUtils.ThreadRejectedExecutionHandler());
    private ConnectManager() {
    }

    private static class SingletonHolder {
        private static final ConnectManager<Object> INSTANCE = new ConnectManager<>();
    }

    @SuppressWarnings("unchecked")
    public static <V> ConnectManager<V> getInstance() {
        return (ConnectManager<V>) SingletonHolder.INSTANCE;
    }

    public void updateConnectedServer(List<RpcProtocol> serviceList) {
        // Now using 2 collections to manage the service info and TCP connections because making the connection is async
        // Once service info is updated on ZK, will trigger this function
        // Actually client should only care about the service it is using
        if (serviceList != null && serviceList.size() > 0) {
            // Update local server nodes cache
            Set<RpcProtocol> serviceSet =Sets.newHashSetWithExpectedSize(serviceList.size());
            serviceSet.addAll(serviceList);
            // Add new server info
            serviceSet.stream().filter(rpcProtocol -> !rpcProtocolSet.contains(rpcProtocol)).forEach(this::connectServerNode);
            // Close and remove invalid server nodes
            rpcProtocolSet.stream().filter(rpcProtocol -> !serviceSet.contains(rpcProtocol)).forEach(this::removeNode);
        } else {
            // No available service
            log.error("No available service!");
            rpcProtocolSet.forEach(this::removeNode);
        }
    }

    private void connectServerNode(RpcProtocol rpcProtocol) {
        if (rpcProtocol.getServiceInfoList() == null || rpcProtocol.getServiceInfoList().isEmpty()) {
            log.info("No service on node, host: {}, port: {}", rpcProtocol.getHost(), rpcProtocol.getPort());
            return;
        }
        rpcProtocolSet.add(rpcProtocol);
        log.info("New service node, host: {}, port: {}", rpcProtocol.getHost(), rpcProtocol.getPort());
        for (RpcServiceInfo serviceProtocol : rpcProtocol.getServiceInfoList()) {
            log.info("New service info, name: {}, version: {}", serviceProtocol.getServiceName(), serviceProtocol.getVersion());
        }
        final InetSocketAddress remotePeer = new InetSocketAddress(rpcProtocol.getHost(), rpcProtocol.getPort());
        threadPoolExecutor.submit(() -> {
            Bootstrap b = new Bootstrap();
            b.group(eventLoopGroup)
                    .channel(NioSocketChannel.class)
                    .handler(new RpcClientInitializer());

            ChannelFuture channelFuture = b.connect(remotePeer);
            channelFuture.addListener((ChannelFutureListener) channelFuture1 -> {
                if (channelFuture1.isSuccess()) {
                    log.info("Successfully connect to remote server, remote peer = " + remotePeer);
                    @SuppressWarnings("unchecked")
                    RpcClientHandler<V> handler = channelFuture1.channel().pipeline().get(RpcClientHandler.class);
                    connectedServerNodes.put(rpcProtocol, handler);
                    handler.setProtocol(rpcProtocol);
                    signalAvailableHandler();
                } else {
                    log.error("Can not connect to remote server, remote peer = " + remotePeer);
                }
            });
        });
    }

    private void signalAvailableHandler() {
        lock.lock();
        try {
            connected.signalAll();
        } finally {
            lock.unlock();
        }
    }

    private boolean waitingForHandler() throws InterruptedException {
        lock.lock();
        try {
            log.warn("Waiting for available service");
            long waitTimeout = 5000;
            return connected.await(waitTimeout, TimeUnit.MILLISECONDS);
        } finally {
            lock.unlock();
        }
    }

    public  RpcClientHandler<V> chooseHandler(String serviceKey) throws Exception {
        int size = connectedServerNodes.values().size();
        while (isRunning && size <= 0) {
            try {
                waitingForHandler();
                size = connectedServerNodes.values().size();
            } catch (InterruptedException e) {
                log.error("Waiting for available service is interrupted!", e);
            }
        }
        RpcProtocol rpcProtocol = loadBalance.route(serviceKey, connectedServerNodes);
        RpcClientHandler<V> handler = connectedServerNodes.get(rpcProtocol);
        if (handler != null) {
            return handler;
        } else {
            throw new Exception("Can not get available connection");
        }
    }

    public void removeHandler(RpcProtocol rpcProtocol) {
        rpcProtocolSet.remove(rpcProtocol);
        connectedServerNodes.remove(rpcProtocol);
        log.info("Remove one connection, host: {}, port: {}", rpcProtocol.getHost(), rpcProtocol.getPort());
    }

    public void stop() {
        isRunning = false;
        for (RpcProtocol rpcProtocol : rpcProtocolSet) {
            removeNode(rpcProtocol);
        }
        rpcProtocolSet.forEach(this::removeNode);
        signalAvailableHandler();
        threadPoolExecutor.shutdown();
        eventLoopGroup.shutdownGracefully();
    }

    private void removeNode(RpcProtocol rpcProtocol){
        RpcClientHandler<V> handler = connectedServerNodes.get(rpcProtocol);
        if (handler != null) {
            handler.close();
        }
        connectedServerNodes.remove(rpcProtocol);
        rpcProtocolSet.remove(rpcProtocol);
    }
}
